# -*-coding: Utf-8 -*-
# @File : BloomFilter .py
# author: Chimengmeng
# blog_url : https://www.cnblogs.com/dream-ze/
# Time：2023/10/7
import hashlib
import math
from . import settings


# 创建布隆过滤器
class BloomFilter:
    def __init__(self, redis_conn):
        # 初始化 Redis 链接对象
        self.redis_conn = redis_conn
        # 初始化 过滤器中元素的数量
        self.num_elements = settings.NUM_ELEMENTS
        # 初始化 期望的误报率
        self.false_positive_rate = settings.FALSE_POSITIVE_RATE

        # 计算位数组的大小
        self.bit_array_size = self.calculate_bit_array_size()
        # 计算哈希函数的数量
        self.num_hashes = self.calculate_num_hashes()

    def calculate_bit_array_size(self):
        # m = -(n * ln(p)) / (ln(2)^2)
        # m 是所需要的位数。n 是过滤器中元素的数量。p 是期望的误报率。
        return int(-(self.num_elements * math.log(self.false_positive_rate)) / (math.log(2) ** 2))

    def calculate_num_hashes(self):
        # k = (m / n) * ln(2)
        return int((self.bit_array_size / self.num_elements) * math.log(2))

    def add(self, element):
        # 使用多个哈希函数计算元素的哈希值
        hashes = [hashlib.sha256(f"{element}{i}".encode()).hexdigest() for i in range(self.num_hashes)]

        # 将位数组中对应位置设置为1
        for hash_value in hashes:
            index = int(hash_value, 16) % self.bit_array_size
            self.redis_conn.setbit("bloom_filter", index, 1)

    def contains(self, element):
        # 使用多个哈希函数计算元素的哈希值
        hashes = [hashlib.sha256(f"{element}{i}".encode()).hexdigest() for i in range(self.num_hashes)]

        # 检查位数组中对应位置是否都为1
        for hash_value in hashes:
            index = int(hash_value, 16) % self.bit_array_size
            if not self.redis_conn.getbit("bloom_filter", index):
                return False

        return True
